package com.allm.ai.common.strategy.impl;

import com.allm.ai.common.dto.ApiRequest;
import com.allm.ai.common.dto.StreamingOutput;
import com.allm.ai.common.entity.AiModel;
import com.allm.ai.common.strategy.AiModelStrategy;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 豆包模型策略实现 - 基于火山引擎官方ArkService
 * 参考官方示例: https://www.volcengine.com/docs/82379/1449737
 */
@Slf4j
@Component
public class DouBaoStrategy implements AiModelStrategy {
    
    @Override
    public String callApi(AiModel model, ApiRequest request) throws Exception {
        log.info("调用豆包API (同步)，模型: {}", model.getName());
        
        ArkService arkService = null;
        try {
            // 创建ArkService实例
            arkService = ArkService.builder()
                .apiKey(model.getApiKey())
                .timeout(Duration.ofMinutes(30)) // 深度思考模型需要较长时间
                .build();
            
            // 构建消息列表
            List<ChatMessage> chatMessages = buildChatMessages(request);
            
            // 构建请求
            ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(model.getRequestModelName())
                .messages(chatMessages)
                .build();
            
            // 发送请求
            var response = arkService.createChatCompletion(chatCompletionRequest);
            
            // 处理响应
            if (response.getChoices() != null && !response.getChoices().isEmpty()) {
                var choice = response.getChoices().get(0);
                var message = choice.getMessage();
                
                // 检查是否有推理内容（thinking模型）
                if (isReasoningModel(model) && message.getReasoningContent() != null) {
                    log.info("推理内容: {}", message.getReasoningContent());
                }
                
                return message.getContent().toString();
            }
            
            throw new Exception("API响应中没有有效内容");
            
        } catch (Exception e) {
            log.error("豆包API调用失败", e);
            throw new Exception("豆包API调用失败: " + e.getMessage(), e);
        } finally {
            if (arkService != null) {
                arkService.shutdownExecutor();
            }
        }
    }
    
    @Override
    public Flux<StreamingOutput> callApiStream(AiModel model, ApiRequest request) {
        log.info("调用豆包流式API，模型: {}", model.getName());
        
        return Flux.create(sink -> {
            ArkService arkService = null;
            try {
                // 创建ArkService实例
                arkService = ArkService.builder()
                    .apiKey(model.getApiKey())
                    .timeout(Duration.ofMinutes(30))
                    .build();
                
                // 构建消息列表
                List<ChatMessage> chatMessages = buildChatMessages(request);
                
                // 构建流式请求
                ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model(model.getRequestModelName())
                    .messages(chatMessages)
                    .build();
                
                // 发送流式请求
                arkService.streamChatCompletion(chatCompletionRequest)
                    .forEach(response -> {
                        try {
                            if (response.getChoices() != null && !response.getChoices().isEmpty()) {
                                var choice = response.getChoices().get(0);
                                var message = choice.getMessage();
                                
                                if (message != null) {
                                    // 处理推理内容（thinking模型）
                                    if (isReasoningModel(model) && message.getReasoningContent() != null) {
                                        String reasoningContent = message.getReasoningContent();
                                        if (!reasoningContent.trim().isEmpty()) {
                                            log.info("推理内容: {}", reasoningContent);
                                            sink.next(new StreamingOutput(StreamingOutput.THINKING, reasoningContent));
                                        }
                                    }
                                    
                                    // 处理常规内容
                                    if (message.getContent() != null && !message.getContent().toString().trim().isEmpty()) {
                                        log.info("消息内容: {}", message.getContent());
                                        sink.next(new StreamingOutput(StreamingOutput.CONTENT, message.getContent().toString()));
                                    }
                                }
                                
                                // 检查是否完成
                                if (choice.getFinishReason() != null) {
                                    log.info("流式调用完成，原因: {}", choice.getFinishReason());
                                    sink.next(new StreamingOutput(StreamingOutput.DONE, "DONE"));
                                    sink.complete();
                                }
                            }
                        } catch (Exception e) {
                            log.error("处理流式响应时出错", e);
                            sink.error(e);
                        }
                    });
                
            } catch (Exception e) {
                log.error("豆包流式API调用失败", e);
                sink.error(e);
            } finally {
                if (arkService != null) {
                    arkService.shutdownExecutor();
                }
            }
        });
    }
    
    @Override
    public boolean supports(String serviceProvider) {
        return "火山引擎".equals(serviceProvider) || 
               "volcengine".equalsIgnoreCase(serviceProvider) ||
               "volc".equalsIgnoreCase(serviceProvider);
    }
    
    @Override
    public boolean supportsStreaming() {
        return true;
    }
    
    /**
     * 构建聊天消息列表
     */
    private List<ChatMessage> buildChatMessages(ApiRequest request) {
        List<ChatMessage> chatMessages = new ArrayList<>();
        
        for (Map<String, Object> msg : request.getMessages()) {
            String role = (String) msg.get("role");
            String content = (String) msg.get("content");
            
            if (content == null || content.trim().isEmpty()) {
                continue;
            }
            
            ChatMessageRole chatRole;
            switch (role != null ? role.toLowerCase() : "user") {
                case "system":
                    chatRole = ChatMessageRole.SYSTEM;
                    break;
                case "user":
                    chatRole = ChatMessageRole.USER;
                    break;
                case "assistant":
                    chatRole = ChatMessageRole.ASSISTANT;
                    break;
                default:
                    log.warn("未知的消息角色: {}, 作为用户消息处理", role);
                    chatRole = ChatMessageRole.USER;
            }
            
            ChatMessage chatMessage = ChatMessage.builder()
                .role(chatRole)
                .content(content)
                .build();
            
            chatMessages.add(chatMessage);
        }
        
        if (chatMessages.isEmpty()) {
            throw new IllegalArgumentException("消息列表不能为空");
        }
        
        return chatMessages;
    }
}